# -*- coding: utf-8 -*-
# @日期    : 2021/11/28 15:37
# @作者  : 万方名
# @FileName: my_rank.py

import numpy as np
import pandas as pd
import xgboost as xgb

from sklearn.feature_extraction.text import TfidfVectorizer

# process data
# dir = '../data/train.csv'
# train_df = pd.read_csv(dir)
# train_df = train_df.loc[:, ['comment_text', 'target']]
# train_df.to_csv('../data/js_rank.csv')

# read data
train_df = pd.read_csv('../data/js_rank.csv')
doc = list(train_df['comment_text'].values)[:1000]

val_df = pd.read_csv('../data/comments_to_score.csv')
print(val_df.head())


print("计算TF-IDF权重")
transformer = TfidfVectorizer()
X_train = transformer.fit_transform(doc)

train = transformer.transform(list(train_df['comment_text'].values)[:1000]).toarray()
train_label = list(train_df['target'].values)[:1000]

test = transformer.transform(list(train_df['comment_text'].values)[1000:2000]).toarray()
test_label = list(train_df['target'].values)[1000:2000]

val = transformer.transform(list(val_df['text'].values)[:1000]).toarray()

# 生成DMatrix
xgb_train = xgb.DMatrix(train, label=train_label)
xgb_test = xgb.DMatrix(test, label=test_label)
xgb_val = xgb.DMatrix(val)

# train model
params = {
    'objective': 'rank:pairwise',  # 学习目标：排序问题
    'gamma': 1.0,
    'min_child_weight': 0.1,  # 叶子结点最小样本权重和，若节点分裂导致叶子结点的样本权重和小于该值则节点不进行分裂
    'eta': 0.1,  # 学习率
    'max_depth': 6  # 决策树分裂的最大深度
}

num_round = 50
watchlist = [(xgb_train, 'train'), (xgb_test, 'test')]

model = xgb.train(params, xgb_train, num_round, watchlist)
model.save_model('../data/model.xgb')

pred = model.predict(xgb_val)
print('预测完成')
"""
data = pd.read_excel('../data/Concrete_Data.xls')

# 对label进行改名
data.rename(columns={'Concrete compressive strength(MPa, megapascals) ': 'label'}, inplace=True)

mask = np.random.rand(len(data)) < 0.8

train = data[mask]
test = data[~mask]
print(f'len of test:{len(test)}')

# 生成DMatrix
xgb_train = xgb.DMatrix(train.iloc[:, :7], label=train.label)
xgb_test = xgb.DMatrix(test.iloc[:, :7], label=test.label)

# train model
params = {
    'objective': 'rank:pairwise',  # 学习目标：排序问题
    'gamma': 1.0,
    'min_child_weight': 0.1,  # 叶子结点最小样本权重和，若节点分裂导致叶子结点的样本权重和小于该值则节点不进行分裂
    'eta': 0.1,  # 学习率
    'max_depth': 6  # 决策树分裂的最大深度
}

num_round = 50
watchlist = [(xgb_train, 'train'), (xgb_test, 'test')]

model = xgb.train(params, xgb_train, num_round, watchlist)
model.save_model('../data/model.xgb')

pred = model.predict(xgb_test)
print(pred)
print(f'len of pred:{len(pred)}')
"""
